#!/home/robot/anaconda3/envs/yolograsp/bin/python3
import os
import open3d as o3d
import numpy as np


def plane_detect(depth, cam_K, verbose=False):
    '''
    args:
        depth: depth img, np.array
    return:
        result: results of plane detection, dict
            including mask of object and normal of table
    '''
    cam_fx, cam_fy, cam_cx, cam_cy = cam_K[0,0], cam_K[1,1], cam_K[0,2], cam_K[1,2]
    norm_scale = 1000.0
    img_width = depth.shape[1]
    img_height = depth.shape[0]
    xmap = np.array([[i for i in range(img_width)] for j in range(img_height)])
    ymap = np.array([[j for i in range(img_width)] for j in range(img_height)])
    
    mask = depth > 0
    choose = mask[:, :].flatten().nonzero()[0]
    depth_masked = depth[:, :].flatten()[choose][:, np.newaxis]
    xmap_masked = xmap[:, :].flatten()[choose][:, np.newaxis]
    ymap_masked = ymap[:, :].flatten()[choose][:, np.newaxis]
    pt2 = depth_masked / norm_scale  # make z coordinate unit be meter
    pt0 = (xmap_masked - cam_cx) * pt2 / cam_fx  # calculate x coordinate under camera axis
    pt1 = (ymap_masked - cam_cy) * pt2 / cam_fy  # calculate y coordinate under camera axis
    points = np.concatenate((pt0, pt1, pt2), axis=1)
    
    pcd = o3d.geometry.PointCloud()
    pcd.points = o3d.utility.Vector3dVector(points)
    plane_model, inliers = pcd.segment_plane(distance_threshold=0.01, ransac_n=4, num_iterations=128)
    [a, b, c, d] = plane_model

    obj_mask = np.ones_like(mask)
    inliers_np = np.array(inliers)
    inliers_index_x = xmap_masked[inliers_np].squeeze()
    print(inliers_index_x)
    inliers_index_y = ymap_masked[inliers_np].squeeze()
    obj_mask[inliers_index_y, inliers_index_x] = 0

    if verbose:
        print(f"Plane equation: {a:.2f}x + {b:.2f}y + {c:.2f}z + {d:.2f} = 0")
        print(f"||(a,b,c)||_2 = {np.linalg.norm(np.array([a, b, c])): .2f}")
        import matplotlib.pyplot as plt
        plt.imshow(obj_mask)
        plt.show()
        inlier_cloud = pcd.select_by_index(inliers)
        inlier_cloud.paint_uniform_color([1.0, 0, 0])
        outlier_cloud = pcd.select_by_index(inliers, invert=True)
        outlier_cloud.paint_uniform_color([0, 1.0, 0])
        coord = o3d.geometry.TriangleMesh.create_coordinate_frame(size=0.1, origin=[0, 0, 0])
        o3d.visualization.draw_geometries([inlier_cloud, outlier_cloud, coord], point_show_normal=True)

    result = dict()
    result['normal'] = np.array([a,b,c])
    if c > 0:
        result['normal'] *= -1
    result['mask'] = obj_mask

    return result

if __name__ == '__main__':

    def _load_depth(img_path):
        """ Load depth image from img_path. """
        depth_path = img_path + '_depth.tiff'  # for realsense
        import cv2
        depth = cv2.imread(depth_path, -1)
        if len(depth.shape) == 3:
            # This is encoded depth image, let's convert
            # NOTE: RGB is actually BGR in opencv
            depth16 = depth[:, :, 1]*256 + depth[:, :, 2]
            depth16 = np.where(depth16==32001, 0, depth16)
            depth16 = depth16.astype(np.uint16)
        elif len(depth.shape) == 2 and depth.dtype == 'uint16':
            depth16 = depth
        else:
            assert False, '[ Error ]: Unsupported depth type.'
        return depth16

    intrinsics = np.array([[603.1332397460938, 0, 323.1732177734375],
                           [0, 601.6819458007812, 241.5373077392578],
                           [0, 0, 1]])
    dir_path = os.path.dirname(os.path.abspath(__file__))
    print("dir path is : ", dir_path) 
    depth = _load_depth(os.path.join(dir_path, '0000'))
    plane_detect(depth, intrinsics, verbose=True)